// main.cpp
#include <CL/cl.h>
#include <iostream>
#include <vector>
#include <fstream>
#include <sstream>
#include <string>
#include <cstdlib>

#define CHECK_CL(err, msg) do { \
    if ((err) != CL_SUCCESS) { \
        std::cerr << "OpenCL error " << (err) << " at " << msg << std::endl; \
        std::exit(EXIT_FAILURE); \
    } \
} while (0)

std::string loadKernelSource(const std::string& filename) {
    std::string kernel_path = KERNELS_DIR "/" + filename;
    std::ifstream file(kernel_path);
    if (!file.is_open()) {
        throw std::runtime_error("Failed to open kernel file: " + kernel_path);
    }
    return std::string((std::istreambuf_iterator<char>(file)), 
                      std::istreambuf_iterator<char>());
}

int main() {
    const unsigned int N = 1 << 28; // ~268 million elements
    std::cout << "Vector addition of " << N << " elements." << std::endl;
    std::vector<float> A(N), B(N), C(N);
    for (unsigned int i = 0; i < N; ++i) {
        A[i] = static_cast<float>(i) * 0.5f;
        B[i] = static_cast<float>(i) * 2.0f;
    }

    // 1) Platform
    cl_uint num_platforms = 0;
    cl_int err = clGetPlatformIDs(0, nullptr, &num_platforms);
    CHECK_CL(err, "clGetPlatformIDs (count)");
    if (num_platforms == 0) {
        std::cerr << "No OpenCL platforms found." << std::endl;
        return 1;
    }
    std::vector<cl_platform_id> platforms(num_platforms);
    err = clGetPlatformIDs(num_platforms, platforms.data(), nullptr);
    CHECK_CL(err, "clGetPlatformIDs (list)");

    // Choose first platform that has a device
    cl_platform_id platform = platforms[0];

    // 2) Device (try GPU, fallback to CPU)
    cl_device_id device = nullptr;
    cl_uint num_devices = 0;
    // try GPU
    err = clGetDeviceIDs(platform, CL_DEVICE_TYPE_GPU, 1, &device, &num_devices);
    if (err != CL_SUCCESS || num_devices == 0) {
        // fallback to CPU
        err = clGetDeviceIDs(platform, CL_DEVICE_TYPE_CPU, 1, &device, &num_devices);
        CHECK_CL(err, "clGetDeviceIDs (CPU fallback)");
    }

    // Print device name
    {
        char name[256];
        err = clGetDeviceInfo(device, CL_DEVICE_NAME, sizeof(name), name, nullptr);
        CHECK_CL(err, "clGetDeviceInfo (name)");
        std::cout << "Using device: " << name << std::endl;
    }

    // 3) Context & Queue
    cl_context context = clCreateContext(nullptr, 1, &device, nullptr, nullptr, &err);
    CHECK_CL(err, "clCreateContext");
    cl_command_queue queue = clCreateCommandQueueWithProperties(context, device, 0, &err);
    CHECK_CL(err, "clCreateCommandQueueWithProperties");

    // 4) Load & build program
    const std::string kernel_src = loadKernelSource("vector_add.cl");
    const char* src_ptr = kernel_src.c_str();
    size_t src_size = kernel_src.size();
    cl_program program = clCreateProgramWithSource(context, 1, &src_ptr, &src_size, &err);
    CHECK_CL(err, "clCreateProgramWithSource");

    err = clBuildProgram(program, 1, &device, nullptr, nullptr, nullptr);
    if (err != CL_SUCCESS) {
        // print build log
        size_t log_size = 0;
        clGetProgramBuildInfo(program, device, CL_PROGRAM_BUILD_LOG, 0, nullptr, &log_size);
        std::string log;
        log.resize(log_size);
        clGetProgramBuildInfo(program, device, CL_PROGRAM_BUILD_LOG, log_size, &log[0], nullptr);
        std::cerr << "Build log:\n" << log << std::endl;
        CHECK_CL(err, "clBuildProgram");
    }

    cl_kernel kernel = clCreateKernel(program, "vector_add", &err);
    CHECK_CL(err, "clCreateKernel");

    // 5) Buffers
    cl_mem bufA = clCreateBuffer(context, CL_MEM_READ_ONLY | CL_MEM_COPY_HOST_PTR, sizeof(float) * N, A.data(), &err);
    CHECK_CL(err, "clCreateBuffer A");
    cl_mem bufB = clCreateBuffer(context, CL_MEM_READ_ONLY | CL_MEM_COPY_HOST_PTR, sizeof(float) * N, B.data(), &err);
    CHECK_CL(err, "clCreateBuffer B");
    cl_mem bufC = clCreateBuffer(context, CL_MEM_WRITE_ONLY, sizeof(float) * N, nullptr, &err);
    CHECK_CL(err, "clCreateBuffer C");

    // 6) Set kernel args
    err = clSetKernelArg(kernel, 0, sizeof(cl_mem), &bufA);
    CHECK_CL(err, "clSetKernelArg 0");
    err = clSetKernelArg(kernel, 1, sizeof(cl_mem), &bufB);
    CHECK_CL(err, "clSetKernelArg 1");
    err = clSetKernelArg(kernel, 2, sizeof(cl_mem), &bufC);
    CHECK_CL(err, "clSetKernelArg 2");
    err = clSetKernelArg(kernel, 3, sizeof(unsigned int), &N);
    CHECK_CL(err, "clSetKernelArg 3");

    // 7) Enqueue kernel
    size_t global_work_size = N;
    size_t local_work_size = 0; // let runtime decide
    // optional: pick a local size (e.g., 256) if desired:
    // local_work_size = 256;
    if (local_work_size == 0) {
        err = clEnqueueNDRangeKernel(queue, kernel, 1, nullptr, &global_work_size, nullptr, 0, nullptr, nullptr);
    } else {
        err = clEnqueueNDRangeKernel(queue, kernel, 1, nullptr, &global_work_size, &local_work_size, 0, nullptr, nullptr);
    }
    CHECK_CL(err, "clEnqueueNDRangeKernel");

    // 8) Read back result
    err = clEnqueueReadBuffer(queue, bufC, CL_TRUE, 0, sizeof(float) * N, C.data(), 0, nullptr, nullptr);
    CHECK_CL(err, "clEnqueueReadBuffer");

    // 9) Verify
    bool ok = true;
    for (unsigned int i = 0; i < N; ++i) {
        float want = A[i] + B[i];
        if (std::abs(C[i] - want) > 1e-5f) {
            std::cerr << "[Info] Mismatch at " << i << ": got " << C[i] << " want " << want << std::endl;
            ok = false;
            break;
        }
    }
    if (ok) std::cout << "[Info] Success: all results match." << std::endl;

    // 10) Cleanup
    clReleaseMemObject(bufA);
    clReleaseMemObject(bufB);
    clReleaseMemObject(bufC);
    clReleaseKernel(kernel);
    clReleaseProgram(program);
    clReleaseCommandQueue(queue);
    clReleaseContext(context);

    return ok ? 0 : 2;
}
